# -*- coding: utf-8 -*-  
import os
import numpy as np
def remove(file):
    if os.path.exists(file):os.remove(file)
def mkdir(ddir):
    if not os.path.exists(ddir):os.makedirs(ddir)
def rename(scr,des):
    if os.path.exists(scr):
        remove(des)
        os.rename(scr,des)
        remove(scr)
    else:print('rename error:','%s is not exists.'%scr)
def copy(scr,des):
    if os.path.exists(scr):
        open(des,'wb').write(open(scr,'rb').read())
    else:print('copy error:','%s is not exists.'%scr)
import tushare as ts
import datetime as dt
import re
import time
import pandas as pd
def gettoken(file='token.txt'):
    if os.path.exists(file):
        with open(file,'r') as f:
            token = f.readline()
        return(token)
    else:
        print('file not exists!check it')
        return('')
def lasttradeday(rtype='str'):
    yestoday = dt.date.today()-dt.timedelta(days=1)
    weekday = dt.datetime.isoweekday(yestoday)
    if weekday in [1,2,3,4,5]:
        if rtype=='datetime':return(yestoday)
        else:return(yestoday.strftime('%Y%m%d'))
    elif weekday==6:
        tmpday = dt.date.today()-dt.timedelta(days=2)
        if rtype=='datetime':return(tmpday)
        else:return(tmpday.strftime('%Y%m%d'))
    elif weekday==7:
        tmpday = dt.date.today()-dt.timedelta(days=3)
        if rtype=='datetime':return(tmpday)
        else:return(tmpday.strftime('%Y%m%d'))
    else:
        print("There's something wrong in weekday calculators!")
        return
def latestday(rtype='str'):
    if dt.datetime.isoweekday(dt.date.today()) in [1,2,3,4,5] \
       and dt.datetime.now().hour>15:
        if rtype=='datetime':return(dt.date.today())
        else:return(dt.date.today().strftime('%Y%m%d'))
    yestoday = dt.date.today()-dt.timedelta(days=1)
    weekday = dt.datetime.isoweekday(yestoday)
    if weekday in [1,2,3,4,5]:
        if rtype=='datetime':return(yestoday)
        else:return(yestoday.strftime('%Y%m%d'))
    elif weekday==6:
        tmpday = dt.date.today()-dt.timedelta(days=2)
        if rtype=='datetime':return(tmpday)
        else:return(tmpday.strftime('%Y%m%d'))
    elif weekday==7:
        tmpday = dt.date.today()-dt.timedelta(days=3)
        if rtype=='datetime':return(tmpday)
        else:return(tmpday.strftime('%Y%m%d'))
    else:
        print("There's something wrong in weekday calculators!")
        return
def getnum(strs):
    nums=re.sub("\D", "", str(strs))
    return(nums)

def getpro():
    ts.set_token(gettoken())
    pro = ts.pro_api()
    return(pro)
def getstocks(file='stocks.npy',update=False):
    if os.path.exists(file):
        stocks = list(np.load(file))
        nowdate = dt.date.today().strftime('%Y%m%d')
        filetimestamp=os.path.getmtime(file)
        filedate=dt.datetime.fromtimestamp(filetimestamp).strftime('%Y%m%d')
        if nowdate==filedate and update==False:return(stocks)
    else:
        stocks=[]
    pro = getpro()
    try:
        df = pro.stock_basic(exchange='',list_status='L',fields='ts_code')
        df2= pro.stock_basic(exchange='',list_status='D',fields='ts_code')
    except Exception:
        return(stocks)
    new= list(df['ts_code'])
    out = list(df2['ts_code'])
    total =[]
    for i in new:
        if i not in stocks:
            stocks.append(i)
    dels=[]
    for i in stocks:
        if i in out:
            print(i,'had been delisted.')
            stocks.remove(i)
    stocks = np.array(stocks)
    np.save(file,stocks)
    print('stocks file update.')
    return(list(stocks))

##----------mongo--------------##
import pymongo
def mongocon():
    myclient=pymongo.MongoClient("mongodb://localhost:27017/")
    mydb=myclient['stock']
    return(myclient,mydb)
def mgtodf(cur,no_id=True):
    data=list(cur)
    if len(data)<=0:
        return(pd.DataFrame([]))
    df = pd.DataFrame(data,columns=data[0].keys())
    if no_id:del df['_id']
    return(df)
def trade_days(start_date,end_date):#获取交易日历
    con,db=mongocon()
    days=db['trade_cal']
    cur=days.find({"cal_date":{'$gte':start_date,'$lte':end_date}})
    df = mgtodf(cur)
    if df.size>0:return(list(df['cal_date']))
    else:return([])
def writetocsv(file,data):#data is np.array or list
    sstr = ''
    for i in data:
        sstr+=str(i)+','
    with open(file,'a') as f:
        f.write(sstr[:-1]+'\n')
def writetofile(txtfile,strdate):
    with open(txtfile,'a')as f:
        f.write(strdate+'\n')
def getfig(data,name=None):
    data=np.array(data)
    x = [i for i in range(data.size)]
    fig = plt.figure()
    plt.title(name)
    plt.plot(x,data)
    plt.grid()
    return(fig)
def gethkhold(end_date,ratio=1):
    con,db=mongocon()
    hold  = db['hk_hold']
    start_date=(dt.datetime.strptime(end_date,'%Y%m%d')-dt.timedelta(days=30)).strftime('%Y%m%d')
    tradedays=trade_days(start_date,end_date)
    hkhold=[]
    for i in tradedays[::-1]:
        cur = hold.find({"trade_date":i})
        df = mgtodf(cur)
        if df.size>0:
            hkhold1 = list(df[df['ratio']>ratio]['ts_code'])
            for i in hkhold1:
                if 'HK' not in i:
                    hkhold.append(i)
            con.close()
            return(hkhold)
        else:
            print(i,'hk_hold data not exists.')
    if len(hkhold)<=0:
        gethkhold(start_date)
